#!/usr/bin/env python3
# Copyright 2017 Canonical Ltd.  This software is licensed under the
# GNU Affero General Public License version 3 (see the file LICENSE).

"""Atomically write a file.

Reads the byte content from standard-in and writes to the specified file. The
filename is checked against an internal white list. As such it's intended to
be used behind `sudo`.
"""

import argparse
import os
import pipes
import sys

from provisioningserver.utils.fs import atomic_write


DATA_FILES = {
    os.path.join(os.getenv("MAAS_DATA", "/var/lib/maas"), path)
    for path in (
            "dhcpd-interfaces",
            "dhcpd.conf",
            "dhcpd6-interfaces",
            "dhcpd6.conf",
    )
}

CONF_FILES = {
    "/etc/chrony/chrony.conf",
    "/etc/chrony/maas.conf"
}
# For DEVELOPMENT ONLY update config file paths to be prefixed
# with MAAS_ROOT, if defined. Check real and effective UIDs to be super extra
# paranoid (only the latter actually matters).
if os.getuid() != 0 and os.geteuid() != 0:
    root = os.environ.get("MAAS_ROOT")
    if root is not None:
        CONF_FILES = {
            os.path.abspath(os.path.join(root, path.lstrip("/")))
            for path in CONF_FILES
        }

WRITABLE_FILES = DATA_FILES | CONF_FILES


def octal(string):
    """Parse `string` as an octal number."""
    return int(string, 8)


arg_parser = argparse.ArgumentParser(description=__doc__)
arg_parser.add_argument("filename", help="The file to write.")
arg_parser.add_argument("mode", type=octal, help="The octal file mode.")


def main(args, fin):

    # Validate the filename here because using a `choices` argument in the
    # parser results in ugly help and error text.
    if args.filename not in WRITABLE_FILES:
        arg_parser.error(
            "Given filename %s is not in the allowed list. "
            "Choose from: %s." % (
                pipes.quote(args.filename), ", ".join(
                    map(pipes.quote, sorted(WRITABLE_FILES)))))

    # Do not allow "high" bits in the mode, especially setuid and setgid.
    elif args.mode & 0o777 != args.mode:
        arg_parser.error(
            "Given file mode 0o%o is not permitted; only "
            "permission bits may be set." % args.mode)

    # Okay, good to go.
    else:
        atomic_write(
            fin.read(), args.filename, overwrite=True,
            mode=args.mode)


if __name__ == "__main__":
    main(arg_parser.parse_args(), sys.stdin.buffer)
